import os
import sys
import argparse
import cv2
import numpy as np
from PIL import Image
import glob
import csv
import re
from pebble import ProcessPool, ProcessExpired
from multiprocessing import cpu_count
import time

# 支持的图片格式
SUPPORTED_FORMATS = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp']

def is_supported_file(filepath):
    """检查文件是否为支持的图片格式"""
    ext = os.path.splitext(filepath)[1].lower()
    return ext in SUPPORTED_FORMATS

def process_image(filepath, crop_rect, output_dir):
    """处理单张图片：裁剪、转换YUV、计算标准差、保存结果"""
    try:
        # 使用PIL打开图片（兼容更多格式）
        pil_img = Image.open(filepath)
        img = np.array(pil_img)

        # 调整裁剪区域确保在图像范围内
        height, width = img.shape[:2]
        x, y, w, h = crop_rect

        x = max(0, min(x, width - 1))
        y = max(0, min(y, height - 1))
        w = min(w, width - x)
        h = min(h, height - y)

        # 裁剪图像
        cropped = img[y:y+h, x:x+w]

        # 转换为YUV并计算标准差
        yuv = cv2.cvtColor(cropped, cv2.COLOR_RGB2YUV)
        y_std = np.std(yuv[:,:,0])
        u_std = np.std(yuv[:,:,1])
        v_std = np.std(yuv[:,:,2])

        
        # 保存裁剪后的图像
        filename = os.path.basename(filepath)
        name, ext = os.path.splitext(filename)
        output_img_path = os.path.join(output_dir, f"{name}_crop{ext}")
        Image.fromarray(cropped).save(output_img_path)
        
        # 提取时间信息
        file_name = os.path.basename(filepath)
        res = re.search(r'(\d{4}\.\d{2}.\d{2}_\d{2}.\d{2}.\d{2})', file_name)
        if res is not None:
            file_name = res.group(1)
        
        print(f"crop image save to : {output_img_path}")
        return {
            'time': file_name,
            'image': output_img_path,
            'y_std': f"{y_std:.4f}",
            'u_std': f"{u_std:.4f}",
            'v_std': f"{v_std:.4f}",
            'filepath': filepath,
            'error': ''
        }
    except Exception as e:
        file_name = os.path.basename(filepath)
        res = re.search(r'(\d{4}\.\d{2}.\d{2}_\d{2}.\d{2}.\d{2})', file_name)
        if res is not None:
            file_name = res.group(1)
        
        return {
            'time': file_name,
            'image': '',
            'y_std': '0',
            'u_std': '0',
            'v_std': '0',
            'filepath': filepath,
            'error': str(e)
        }


def write_results_to_csv(results, csv_path):
    """将结果写入CSV文件"""
    results.sort(key=lambda x: x['time'])
    max_y_std = max(float(res['y_std']) for res in results)
    with open(csv_path, 'w', newline='', encoding='utf-8') as csvfile:
        csv_writer = csv.writer(csvfile)
        csv_writer.writerow(['minute','time', 'image', 'Y_std', 'U_std', 'V_std', 'attenuation'])
        for i,result in enumerate(results):
            y_std = float(result['y_std'])
            atten_str = f'{y_std/max_y_std:.2f}'
            csv_writer.writerow([
                str(i),
                result['time'],
                result['image'],
                result['y_std'],
                result['u_std'],
                result['v_std'],
                atten_str
            ])
            
            if result['error']:
                print(f"处理失败: {result['filepath']} - {result['error']}")

def main():
    parser = argparse.ArgumentParser(description='图像裁剪与分析工具')
    parser.add_argument('input_path', help='输入图片路径（文件或目录）')
    parser.add_argument('-r', '--rect', type=int, nargs='+', required=True, 
                       help='裁剪区域 [x,y,width,height] 例如: [100,50,200,150]')
    parser.add_argument('-o', '--output', required=True, help='输出目录')
    parser.add_argument('-m','--multiprocess', action='store_true', 
                       help='启用多进程处理')

    args = parser.parse_args()

    # 验证裁剪区域参数
    if len(args.rect) != 4:
        print("错误：裁剪区域必须包含4个参数 [x,y,width,height]")
        return

    # 创建输出目录
    os.makedirs(args.output, exist_ok=True)

    # 处理输入路径
    if os.path.isfile(args.input_path):
        files = [args.input_path] if is_supported_file(args.input_path) else []
    elif os.path.isdir(args.input_path):
        files = []
        for root, dirs, filenames in os.walk(args.input_path):
            for filename in filenames:
                filepath = os.path.join(root, filename)
                if is_supported_file(filepath):
                    files.append(filepath)
    else:
        print("错误：无效的输入路径")
        return

    if not files:
        print("错误：未找到支持的图片文件")
        return
    print(f"找到 {len(files)} 个图片文件")
    results = []
    # 确定是否使用多进程
    if args.multiprocess:
        print("Multi-processing mode.")
        processes_num = min(len(files), cpu_count())
        print(f"Using {processes_num} processes.")
        with ProcessPool(max_workers=processes_num) as pool:
            futures = []
            try:
                # 提交所有转换任务
                for file_path in files:
                    future = pool.schedule(
                        process_image, 
                        args=(file_path, args.rect, args.output)
                    )
                    futures.append(future)
                for future in futures:
                    result = future.result()
                    results.append(result)
            except KeyboardInterrupt:
                print("\n用户中断请求，正在取消所有任务...")
                # 取消所有未完成的任务
                for future in futures:
                    if not future.done():
                        future.cancel()
                # 等待已启动的任务完成
                pool.stop()
                pool.join(timeout=5)
                print("所有任务已终止")
                sys.exit(1)
    else:
        print("Single process mode.")
        for file_path in files:
            results.append(process_image(file_path, args.rect, args.output))
    # 写入CSV文件
    csv_path = os.path.join(args.output, 'crop_analysis.csv')
    write_results_to_csv(results, csv_path)
    
    print(f"analysis result save to: {csv_path}")
    print("done")

if __name__ == "__main__":
    main()
